#include <sys/time.h>
#include <sys/resource.h>
#include <stdio.h>
#include <fstream>
#include <string>
#include <fstream>
#include <armadillo>
#include <qk.h>
#include <io.h>
#include <propagator.h>

//*************************************************************************
void default_parameters(parameters *);
void print_parameters(const parameters& );

void fill_arrays(const parameters& ,
		 complex_matrix_function& Hamiltonian,
		 complex_matrix_function& GFarray);
void micro_time(parameters& params,const arma::vec& tdiag, arma::vec& tau);
//*************************************************************************

//************************************************************************
int main(int argc, char **argv){
  using namespace std;
  using namespace arma;

  // INITIALIZE TO DEFAULT PARAMETERS 
  parameters params;
  default_parameters(&params);

  //
  //     SET PARAMETERS BY READING THE INPUT COMMAND LINE
  //
  //   Argument order: see usage statement below
  //                   that can be produced by "help" as an argument

  
  if (argc>1)
    {
      string arg1;
      arg1 = argv[1];
      if (arg1=="help")
	{
	  cout << "\n usage: qk [input h5 file] [memwidth] \n\n";
	  return 0;
	}
    }

  string inputfilename("input.h5");
  int iarg=1;
  if (argc>iarg)
    inputfilename = argv[iarg];
  iarg++;
  if (argc>iarg)
    params.memwidth    = atof(argv[iarg]);

  // STORE THE CALLING COMMAND IN FILE lastcommand 
  ofstream cmdfile;
  cmdfile.open("lastcommand.sh");
  
  for (int i=0; i < argc ; i++)
    {
      cmdfile << argv[i] << " ";
    }
  cmdfile << "\n";
  cmdfile.close();


  // TIME ARRAYS: params.ntmem is changed here
  // Hamiltonian and G.F. arrays 
  //	
  //	Hamiltonian:
  //	This is STL vectors with each member corresponding to a k point.
  //	Each member of the vector is a nband x nband matrix which is
  //	currently armadillo cx_mat
  vec tdiag, tau,Efield;
  mat kpoints;
  mat EnergyBands;
  complex_matrix_function IntHamiltonian,rho0;
  io::read_input_hdf5_file("input.h5",tdiag,Efield,kpoints,EnergyBands,IntHamiltonian,rho0);
  params.ntdiag  = tdiag.n_elem;
  params.tstart  = tdiag(0);
  params.tend    = tdiag(params.ntdiag-1);
  // MEMORY TIME
  micro_time(params,tdiag, tau);

  params.num_mesh_points=kpoints.n_rows;
  params.deltakx = kpoints(1,0)-kpoints(0,0);
  if (params.deltakx==0)
    {
      cerr<<"\n\nERROR: main() params.deltakx is 0\n\n";
    }
  params.num_mesh_points_x = 1;
  {
    int i=1;
    while ( fabs(kpoints(i,1)-kpoints(i-1,1)) < 1.0e-10 && i < params.num_mesh_points )
      {
	i++;
	if (i >= params.num_mesh_points)
	  break;
      }
    params.num_mesh_points_x = i;
  }
    
  // PRINTOUT: parameters
  print_parameters(params);
  
  // FILENAMES
  //    Set the file names for I/O operations during calculations
  //    and initialize the HDF5 files for chunked data transfer
  h5files_container files; 
  io::set_filenames(files);
  io::disk_storage_initialization(files, params);
  
  //	GFarray:
  //	This is also STL vector that has num kpoints x num tau points many members
  //	Each member is a nband x nband matrix
  unsigned long ntot=params.num_mesh_points * params.ntmem;
  complex_matrix_function GFarray(ntot);
  for (int i=0; i<ntot;i++)
    GFarray[i].zeros(params.num_bands,params.num_bands);
  for (int ik=0; ik<params.num_mesh_points;ik++)
    GFarray[ik*params.ntmem] = rho0[ik];


  // CALL TO PROPAGATOR 
  propagation::propagator( params, tdiag, tau, Efield, EnergyBands,IntHamiltonian, files,GFarray);

  // FINISH
  cout<<"\nFINISHED\n";
  return 0;

}


//*************************************************************************
void default_parameters(parameters *p){
  p->tstart    =  0.0;
  p->tend      =  8.0;
  p->memwidth  =  0.1;
  p->ntdiag    =  20;
  p->ntmem     =  10;
  
  p->num_mesh_points_x = 3;
  p->deltakx           = 0.3;

  p->num_bands =  2;
  p->num_mesh_points = 3;

  // Control
  p->tolerance = 0.1;
  p->maxiter   = 1;

};
void print_parameters(const parameters& params){
  printf("tstart            %14.4f\n",params.tstart);
  printf("tend              %14.4f\n",params.tend);
  printf("memwidth          %14.4f\n",params.memwidth);
  printf("tolerance         %14.4f\n",params.tolerance);
  printf("ntdiag            %14d\n",params.ntdiag);
  printf("ntmem             %14d\n",params.ntmem);
  printf("deltakx           %14.4f\n",params.deltakx);
  printf("num_bands         %14d\n",params.num_bands);
  printf("num_mesh_points   %14d\n",params.num_mesh_points);
  printf("num_mesh_points_x %14d\n",params.num_mesh_points_x);
  printf("tolerance         %14d\n",params.tolerance);
  printf("maxiter           %14d\n",params.maxiter);

}
void fill_arrays(const parameters& params,
		 complex_matrix_function& Hamiltonian,
		 complex_matrix_function& GFarray)
{
  using namespace arma;
  using namespace std;
  //----------------------------------------------------------------------
  // FILL HAMILTONIAN ARRAY AND INITIAL GF ARRAY
  // TODO: Filling this array should be moved to the io file	
  
  double aa=0.23;
  for (int ik=0; ik < params.num_mesh_points; ik++)
    {
      int itau;
      for (itau=0;itau<params.ntmem; itau++)
	GFarray[itau + ik*params.ntmem].zeros(params.num_bands,params.num_bands);

      // The green functiona at tau = 0 
      itau=0;
      GFarray[itau + ik*params.ntmem](0,0) = aa;
      GFarray[itau + ik*params.ntmem](1,1) = aa;
      GFarray[itau + ik*params.ntmem](0,1) = aa;
      GFarray[itau + ik*params.ntmem](1,0) = aa;
      Hamiltonian[ik].zeros(params.num_bands,params.num_bands);
      Hamiltonian[ik](0,0) =  3.4 * pi;
      Hamiltonian[ik](1,1) = -1 * pi;
      Hamiltonian[ik](0,1) =  sqrt(8.0)* pi ;
      Hamiltonian[ik](1,0) =  sqrt(8.0)* pi ;
    }

  // Write Hamiltonian and initial DM to disk 
  mat r;
  r=real(Hamiltonian[0]);
  r.save("hamiltonian_real.dat",raw_ascii);
  r=imag(Hamiltonian[0]);
  r.save("hamiltonian_imag.dat",raw_ascii);  
  r = real(GFarray[0]);
  r.save("rho.dat",raw_ascii);

}
void micro_time(parameters& params,const arma::vec& tdiag, arma::vec& tau)
{
  using namespace std;
  using namespace arma;
  // Time arrays 
  //   -------  <=  mem  =>
  // width 
  // --------. . . . . .tau5
  // ------. . . . . .  tau4
  // ----. . . . . .    tau3
  // --. . . . . .      tau2
  // . . . . . .        tau1
  // t1  t3  t5
  //
  // t1, t2 t3 ... are give in tdiag

  // identify the index where t = memory width 
  int itau=0;
  while (itau < params.ntdiag-1 & tdiag(itau)-tdiag(0)<params.memwidth)
    itau++;
  // Make sure itau does not exceed the last index in tdiag array 
  itau = min(params.ntdiag-1,max(1,itau));

  // Copy tdiag form 0...itau into tau subtracting tdiag(0) from it
  // to allow for any offset, as was done above in the while loop as well
  
  tau             = tdiag.subvec(0,itau)-tdiag(0);
  params.ntmem    = tau.n_elem;
  params.memwidth = tau(tau.n_elem-1);

  tau.save("tau.dat",raw_ascii);
  tdiag.save("time.dat",raw_ascii);

}
